# 处理多个回调

from typing import Any, Dict, List, Union
from langchain.agents import AgentType, initialize_agent
from langchain_community.agent_toolkits.load_tools import load_tools
from langchain.callbacks.base import BaseCallbackHandler
from langchain_core.agents import AgentAction
from langchain_ollama.llms import OllamaLLM

# 自定义回调处理函数
class MyCustomHandlerOne(BaseCallbackHandler):
    def on_llm_start(self, serialized, prompts, **kwargs):
        print(f"on_llm_start {serialized['name']}")

    def on_llm_new_token(self, token: str, **kwargs: Any) -> Any:
        print(f"on_new_token {token}")

    def on_llm_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> Any:
        pass

    def on_chain_start(self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs) -> Any:
        # 添加检查，避免 serialized 为 None 时出错
        if serialized is not None and 'name' in serialized:
            print(f"on_chain_start {serialized['name']}")

    def on_tool_start(self, serialized: Dict[str, Any], input_str: str, **kwargs: Any) -> Any:
        print(f"on_tool_start {serialized['name']}")

    def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
        print(f"on_agent_action {action}")

class MyCustomHandlerTwo(BaseCallbackHandler):
    def on_llm_start(self, serialized, prompts, **kwargs):
        print(f"on_llm_start (i am the second handler!) {serialized['name']}")

# 实例化处理器
handler1 = MyCustomHandlerOne()
handler2 = MyCustomHandlerTwo()
llm = OllamaLLM(model="deepseek-r1:7b", callbacks = [handler2])

# 修正传入 load_tools 的参数
tools = load_tools(["llm-math"], llm = llm)

# 设置 handle_parsing_errors=True 以处理输出解析错误
# 采用零样本推理代理
agent = initialize_agent(tools, llm, agent = AgentType.ZERO_SHOT_REACT_DESCRIPTION)

# 使用invoke方法替代run方法,run已经弃用
result = agent.invoke({"input": "为什么影子之间会有吸引力？"}, callbacks=[handler1])
print(result)

'''理论流程：
1,代理启动，触发 on_chain_start 回调。
2,语言模型启动，触发 on_llm_start 回调（handler2 和 handler1 都会触发）。
3,语言模型生成 token，触发 on_llm_new_token 回调。
4,代理执行工具（如 llm-math），触发 on_tool_start 回调。
5,代理执行动作，触发 on_agent_action 回调。
最终输出代理的运行结果。
'''


'''此代码暂未调通！
    V25-3-4 19:35: 运行时间过长！且重复输出：on_llm_start (i am the second handler!) OllamaLLM
    原因可能是llm的输出不符合agent的期望
'''
   